use polars::prelude::*;
use pyo3::prelude::*;
use pyo3_polars::PyDataFrame;
use std::collections::{HashMap, HashSet};
mod record;
use record::*;
const BUY: &str = "buy";
const SELL: &str = "sell";
const BUY_SHORT: &str = "buy_short";
const SELL_SHORT: &str = "sell_short";

impl BackTest {
    // 获取详细的交易记录
    fn get_trade_details(&self, strategy: String, params: Option<String>) -> PyResult<Option<Vec<f64>>> {
        match params {
            None => {
                // 单次回测:生成完整报告
                self.record.write_all_to_csv(&strategy).unwrap();
                self.record
                    .generate_html_report(&strategy, self.init_cash)
                    .expect("Failed to generate HTML report");
                Ok(None)
            }
            Some(_) => {
                // 批量回测:计算指标并返回
                let metrics = self.record.calculate_batch_metrics(self.init_cash).unwrap();
                Ok(Some(vec![
                    self.init_cash,
                    metrics.final_assets,
                    metrics.total_profit_loss,
                    metrics.total_return,
                    metrics.annualized_return,
                    metrics.max_drawdown,
                    metrics.total_trades as f64,
                    metrics.trading_days as f64,
                ]))
            }
        }
    }
}
#[pyclass]
pub struct BackTest {
    // 内部使用
    position_to_trade: HashMap<String, String>, //交易代码到数据代码的映射 "CFFEX_IF00" -> "CFFEX_IF00_60"
    trade_data: HashMap<String, DataFrame>, // 行情合约数据
    record: Record, // 记录对象
    init_cash: f64, // 初始资金
    symbol_index: HashMap<String, usize>, // 行情代码到索引的映射

    // 对外暴露的字段
    #[pyo3(get)]
    cash: f64, //可用资金
    #[pyo3(get)]
    fee: f64, //手续费率
    #[pyo3(get)]
    deposit: HashMap<String, f64>, //保证金率
    #[pyo3(get)]
    multiplier: HashMap<String, f64>, //品种单位
    #[pyo3(get)]
    today: String, // 当天日期
    #[pyo3(get)]
    timestamps: i64, // 当天时间戳
    #[pyo3(get)]
    stock_value: HashMap<String, f64>, //个股市值/保证金
    #[pyo3(get)]
    market_value: f64, //总市值
    #[pyo3(get)]
    stock_profit: HashMap<String, f64>, //个股浮盈
    #[pyo3(get)]
    profit: f64, //账户盈亏
    #[pyo3(get)]
    long_position: HashMap<String, u32>, //多头持仓
    #[pyo3(get)]
    short_position: HashMap<String, u32>, //空头持仓
    #[pyo3(get)]
    long_average_price: HashMap<String, f64>, //多头持仓均价
    #[pyo3(get)]
    short_average_price: HashMap<String, f64>, //空头持仓均价
}

impl BackTest {
    // 每天盘后更新账户信息,累加所有个股市值和浮盈
    fn update_account(&mut self) {
        self.market_value = self.stock_value.values().sum();
        self.profit = self.stock_profit.values().sum();
        // 期货:动态权益 = 可用资金 + 浮动盈亏 + 保证金
        self.record.account_log.push(AccountRecord {
            date: self.today.clone(),
            available_cash: self.cash,
            market_value: self.market_value,
            total_assets: self.cash + self.market_value + self.profit,
            total_profit_loss: self.profit,
        })
    }

    // 每个bar更新账户信息,累加所有个股市值和浮盈
    fn update_account_each_bar(&mut self) {
        self.market_value = self.stock_value.values().sum();
        self.profit = self.stock_profit.values().sum();
        // 期货:动态权益 = 可用资金 + 浮动盈亏 + 保证金
        self.record.account_log_each_bar.push(AccountRecord {
            date: self.today.clone(),
            available_cash: self.cash,
            market_value: self.market_value,
            total_assets: self.cash + self.market_value + self.profit,
            total_profit_loss: self.profit,
        })
    }
    // 用来更新品种持仓信息,当个股有持仓时候更新
    fn update_position(&mut self) {
        // 只处理有实际持仓的品种
        let codes_with_position: HashSet<String> = self
            .long_position
            .iter()
            .chain(self.short_position.iter())
            .filter(|&(_, &qty)| qty > 0) // 先过滤掉持仓为0的
            .map(|(code, _)| code.clone()) // 只克隆有持仓的品种
            .collect();
        for code in &codes_with_position {
            let buy_qty = *self.long_position.get(code).unwrap_or(&0);
            let sell_qty = *self.short_position.get(code).unwrap_or(&0);

            let price = self.get_price_binary(code).unwrap_or(0.0);
            let multiplier = *self.multiplier.get(code).unwrap_or(&0.0);
            let deposit_rate = *self.deposit.get(code).unwrap_or(&0.0);

            // 计算多头
            let (long_pnl, long_value) = if buy_qty > 0 {
                let avg_cost = *self.long_average_price.get(code).unwrap_or(&0.0);
                let pnl = (price - avg_cost) * (buy_qty as f64) * multiplier;
                let value = (buy_qty as f64) * price * multiplier * deposit_rate;
                (pnl, value)
            } else {
                (0.0, 0.0)
            };
            // 计算空头
            let (short_pnl, short_value) = if sell_qty > 0 {
                let avg_cost = *self.short_average_price.get(code).unwrap_or(&0.0);
                let pnl = (avg_cost - price) * (sell_qty as f64) * multiplier;
                let value = (sell_qty as f64) * price * multiplier * deposit_rate;
                (pnl, value)
            } else {
                (0.0, 0.0)
            };
            // 更新记录
            self.stock_value
                .insert(code.clone(), long_value + short_value);
            self.stock_profit.insert(code.clone(), long_pnl + short_pnl);
            // 记录持仓信息,先不启用
            // let position_record = PositionRecord {
            //     date: self.today.clone(),
            //     symbol: code.clone(),
            //     long_quantity: buy_qty,
            //     long_average_cost: *self.long_average_price.get(code).unwrap_or(&0.0),
            //     long_unrealized_pnl: long_pnl,
            //     long_market_value: long_value,
            //     short_quantity: sell_qty,
            //     short_average_cost: *self.short_average_price.get(code).unwrap_or(&0.0),
            //     short_unrealized_pnl: short_pnl,
            //     short_market_value: short_value,
            // };
            // self.record.position_log.push(position_record);
        }
    }

    // 更新交易明细,每次执行交易动作时候更新
    fn update_trade(&mut self, symbol: String, trade_type: &'static str, volume: u32, price: f64) {
        self.record.trade_log.push(TradeRecord {
            date: self.today.clone(),
            symbol,
            trade_type,
            volume,
            price,
        })
    }
    // 二分法提取价格获取逻辑
    fn get_price_binary(&self, code: &str) -> Option<f64> {
        let df = self.trade_data.get(self.position_to_trade.get(code)?)?;
        let timestamps_series = df.column("timestamps").ok()?;
        let close_series = df.column("close").ok()?;
        let timestamps = timestamps_series.i64().ok()?;
        let closes = close_series.f64().ok()?;
        // 使用二分查找
        let mut left = 0;
        let mut right = timestamps.len().saturating_sub(1);
        let mut result = None;
        while left <= right {
            let mid = left + (right - left) / 2;
            if let (Some(ts), Some(close)) = (timestamps.get(mid), closes.get(mid)) {
                if ts <= self.timestamps {
                    result = Some(close);
                    left = mid + 1;
                } else {
                    right = mid - 1;
                }
            } else {
                break;
            }
        }
        result
    }
    // 用索引法提取最新价,当单个品种数据到百万行时,二分法效率较低,用索引法效率较高
    fn get_price_indexed(&mut self, code: &str) -> Option<f64> {
        let trade_code = self.position_to_trade.get(code)?;
        let df = self.trade_data.get(trade_code)?;
        let index = self.symbol_index.get_mut(trade_code)?;

        let timestamps = df.column("timestamps").ok()?.i64().ok()?;
        let closes = df.column("close").ok()?.f64().ok()?;

        let mut price = None;

        while *index < timestamps.len() {
            let ts = timestamps.get(*index)?;
            let close = closes.get(*index)?;

            if ts < self.timestamps {
                price = Some(close);
                *index += 1;
            }
            else if ts == self.timestamps {
                price = Some(close);
                break;
            }
            else {
                // ts > self.timestamps,返回前一个价格
                if *index > 0 {
                    price = closes.get(*index - 1);  // 返回上一个价格
                }
                break;
            }
        }
        price
    }
}

fn extract_position_code(trade_code: &str) -> Option<String> {
    let parts: Vec<&str> = trade_code.split('_').collect();
    if parts.len() >= 2 {
        Some(format!("{}_{}", parts[0], parts[1])) // "CFFEX_IF00"
    } else {
        None
    }
}

#[pymethods]
impl BackTest {
    #[new]
    fn new(
        cash: f64,
        fee: f64,
        deposit: HashMap<String, f64>,
        multiplier: HashMap<String, f64>,
        trade_data: HashMap<String, PyDataFrame>,
    ) -> PyResult<Self> {
        // 建立映射关系
        let mut position_to_trade = HashMap::with_capacity(trade_data.len());
        let mut symbol_index = HashMap::with_capacity(trade_data.len());
        for trade_code in trade_data.keys() {
            if let Some(position_code) = extract_position_code(trade_code) {
                position_to_trade.insert(position_code, trade_code.clone());
                symbol_index.insert(trade_code.clone(), 0);
            }
        }
        // pyo3的dataframe转dataframe
        let mut trade_data_rust = HashMap::with_capacity(trade_data.len());
        for (code, py_df) in trade_data.iter() {
            let df: DataFrame = py_df.clone().into();
            trade_data_rust.insert(code.clone(), df);
        }
        Ok(Self {
            trade_data: trade_data_rust,
            position_to_trade,
            symbol_index,
            cash,
            fee,
            init_cash: cash,
            record: Record {
                trade_log: Vec::new(),
                account_log: Vec::new(),
                account_log_each_bar: Vec::new(),
                position_log: Vec::new(),
            },
            deposit,
            multiplier,
            today: String::new(),
            timestamps: 0,
            stock_value: HashMap::new(),
            market_value: 0.0,
            stock_profit: HashMap::new(),
            profit: 0.0,
            long_position: HashMap::new(),
            short_position: HashMap::new(),
            long_average_price: HashMap::new(),
            short_average_price: HashMap::new(),
        })
    }
    fn buy(&mut self, symbol: String, volume: u32, price: f64) {
        // 1. 参数检查
        if volume <= 0 || price <= 0.0 {
            return;
        }
        // 2. 资金检查
        let multiplier = self.multiplier.get(&symbol).copied().unwrap_or(1.0);
        let deposit = self.deposit.get(&symbol).copied().unwrap_or(1.0);
        let volume_f64 = volume as f64;
        let required_cash = price * volume_f64 * multiplier * deposit;
        let available_cash = self.cash + self.profit;
        if available_cash < required_cash {
            println!(
                "{} 资金不足:需要 {:.2}, 可用 {:.2}",
                self.today, required_cash, available_cash
            );
            return;
        }
        // 3. 更新资金
        let trade_cost = price * volume_f64 * multiplier;

        self.cash -= trade_cost * self.fee; // 手续费

        self.cash -= trade_cost * deposit; // 保证金

        // 4. 更新持仓
        let current_position = self.long_position.get(&symbol).copied().unwrap_or(0);
        let current_avg_price = self.long_average_price.get(&symbol).copied().unwrap_or(0.0);
        let new_position = current_position + volume;
        let new_avg_price = if current_position > 0 {
            // 有旧持仓:计算加权平均成本
            (current_avg_price * current_position as f64 + price * volume_f64) / new_position as f64
        } else {
            // 新开仓:直接使用当前价格
            price
        };
        self.long_position.insert(symbol.clone(), new_position);
        self.long_average_price
            .insert(symbol.clone(), new_avg_price);
        self.update_trade(symbol.clone(), BUY, volume, price);
    }
    /// 平多
    fn sell(&mut self, symbol: String, volume: u32, price: f64) {
        // 1. 参数检查
        if volume <= 0 || price <= 0.0 {
            return;
        }
        // 2. 持仓检查
        let current_position = self.long_position.get(&symbol).copied().unwrap_or(0);
        if current_position == 0 {
            return;
        }
        let current_price = self.long_average_price.get(&symbol).copied().unwrap_or(0.0);
        let multiplier = self.multiplier.get(&symbol).copied().unwrap_or(1.0);
        let deposit = self.deposit.get(&symbol).copied().unwrap_or(1.0);
        let volume_f64 = volume as f64;
        let new_position = if current_position >= volume {
            current_position - volume
        } else {
            0
        };
        if new_position > 0 {
            // 部分平仓
            self.long_position.insert(symbol.clone(), new_position);
            self.cash -= volume_f64 * price * multiplier * self.fee;
            self.cash += volume_f64 * current_price * multiplier * deposit;
            self.cash += (price - current_price) * volume_f64 * multiplier;
        } else {
            self.cash -= current_position as f64 * price * multiplier * self.fee;
            self.cash += current_position as f64 * current_price * multiplier * deposit;
            self.cash += (price - current_price) * current_position as f64 * multiplier;
            // 全部平仓:移除持仓记录
            self.long_position.remove(&symbol);
            self.long_average_price.remove(&symbol);
            self.stock_profit.remove(&symbol);
            self.stock_value.remove(&symbol);
        }
        let actual_volume = if new_position > 0 {
            volume
        } else {
            current_position
        };
        self.update_trade(symbol.clone(), SELL, actual_volume, price);
    }

    /// 开空(卖出开仓)
    fn buy_short(&mut self, symbol: String, volume: u32, price: f64) {
        // 1. 参数检查
        if volume <= 0 || price <= 0.0 {
            return;
        }

        // 2. 资金检查
        let multiplier = self.multiplier.get(&symbol).copied().unwrap_or(1.0);
        let deposit = self.deposit.get(&symbol).copied().unwrap_or(1.0);
        let volume_f64 = volume as f64;
        let required_cash = price * volume_f64 * multiplier * deposit;
        let available_cash = self.cash + self.profit;

        if available_cash < required_cash {
            println!(
                "{} 资金不足:需要 {:.2}, 可用 {:.2}",
                self.today, required_cash, available_cash
            );
            return;
        }

        // 3. 更新资金
        let trade_cost = price * volume_f64 * multiplier;
        self.cash -= trade_cost * self.fee; // 手续费
        self.cash -= trade_cost * deposit; // 保证金

        // 4. 更新空头持仓
        let current_position = self.short_position.get(&symbol).copied().unwrap_or(0);
        let current_avg_price = self
            .short_average_price
            .get(&symbol)
            .copied()
            .unwrap_or(0.0);
        let new_position = current_position + volume;
        let new_avg_price = if current_position > 0 {
            // 有旧持仓:计算加权平均成本
            (current_avg_price * current_position as f64 + price * volume_f64) / new_position as f64
        } else {
            // 新开仓:直接使用当前价格
            price
        };

        self.short_position.insert(symbol.clone(), new_position);
        self.short_average_price
            .insert(symbol.clone(), new_avg_price);
        self.update_trade(symbol.clone(), BUY_SHORT, volume, price);
    }

    /// 平空(买入平仓)
    fn sell_short(&mut self, symbol: String, volume: u32, price: f64) {
        // 1. 参数检查
        if volume <= 0 || price <= 0.0 {
            return;
        }

        // 2. 持仓检查
        let current_position = self.short_position.get(&symbol).copied().unwrap_or(0);
        if current_position == 0 {
            return;
        }
        let current_price = self
            .short_average_price
            .get(&symbol)
            .copied()
            .unwrap_or(0.0);
        let multiplier = self.multiplier.get(&symbol).copied().unwrap_or(1.0);
        let deposit = self.deposit.get(&symbol).copied().unwrap_or(1.0);
        let volume_f64 = volume as f64;
        let new_position = if current_position >= volume {
            current_position - volume
        } else {
            0
        };

        if new_position > 0 {
            // 部分平仓
            self.short_position.insert(symbol.clone(), new_position);
            self.cash -= volume_f64 * price * multiplier * self.fee; // 手续费
            self.cash += volume_f64 * current_price * multiplier * deposit; // 返还保证金
            self.cash += (current_price - price) * volume_f64 * multiplier; // 空头盈利计算
        } else {
            // 全部平仓
            self.cash -= current_position as f64 * price * multiplier * self.fee; // 手续费
            self.cash += current_position as f64 * current_price * multiplier * deposit; // 返还保证金
            self.cash += (current_price - price) * current_position as f64 * multiplier; // 空头盈利计算
            // 移除持仓记录
            self.short_position.remove(&symbol);
            self.short_average_price.remove(&symbol);
            self.stock_profit.remove(&symbol);
            self.stock_value.remove(&symbol);
        }

        let actual_volume = if new_position > 0 {
            volume
        } else {
            current_position
        };
        self.update_trade(symbol.clone(), SELL_SHORT, actual_volume, price);
    }
}

#[pyfunction]
fn start_backtest(
    py: Python,
    python_obj: Py<PyAny>,
    backtest_instance: Py<BackTest>,
    trade_datetime: Vec<String>,
    trade_timestamps: Vec<i64>,
    strategy: String,
) -> PyResult<()> {
    let mut last_date = String::with_capacity(12);
    for (datetime, timestamps) in trade_datetime.iter().zip(trade_timestamps.iter()) {
        let current_date = &datetime[..10];
        let current_hour = &datetime[11..13];
        // 更新日期时间
        {
            let mut instance: PyRefMut<'_, BackTest> = backtest_instance.borrow_mut(py);
            instance.today = datetime.clone();
            instance.timestamps = *timestamps;
        }
        // 调用Python方法
        python_obj.call_method1(py, "on_bar", (timestamps,))?;
        // 每一根k的更新
        {
            backtest_instance.borrow_mut(py).update_position();
            backtest_instance.borrow_mut(py).update_account_each_bar();
        }
        // 尾盘的更新
        if current_hour == "15" && current_date != last_date {
            backtest_instance.borrow_mut(py).update_account();
            last_date = current_date.to_string();
        }
    }
    backtest_instance.borrow(py).get_trade_details(strategy, None).expect("something went wrong");
    Ok(())
}

#[pyfunction]
fn batch_backtest(
    py: Python,
    python_obj: Py<PyAny>,
    backtest_instance: Py<BackTest>,
    trade_datetime: Vec<String>,
    trade_timestamps: Vec<i64>,
    strategy: String,
    params: String,
) -> PyResult<Option<Vec<f64>>> {
    let mut last_date = String::with_capacity(12);
    for (datetime, timestamps) in trade_datetime.iter().zip(trade_timestamps.iter()) {
        let current_date = &datetime[..10];
        let current_hour = &datetime[11..13];
        // 更新日期时间
        {
            let mut instance: PyRefMut<'_, BackTest> = backtest_instance.borrow_mut(py);
            instance.today = datetime.clone();
            instance.timestamps = *timestamps;
        }
        // 调用Python方法
        python_obj.call_method1(py, "on_bar", (timestamps,))?;
        // 每一根k的更新
        {
            backtest_instance.borrow_mut(py).update_position();
            backtest_instance.borrow_mut(py).update_account_each_bar();
        }
        // 尾盘的更新
        if current_hour == "15" && current_date != last_date {
            backtest_instance.borrow_mut(py).update_account();
            last_date = current_date.to_string();
        }
    }
    backtest_instance.borrow(py).get_trade_details(strategy, Some(params))
}
#[pymodule]
fn backtest(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<BackTest>()?;
    m.add_function(wrap_pyfunction!(start_backtest, m)?)?;
    m.add_function(wrap_pyfunction!(batch_backtest, m)?)?;
    Ok(())
}
